package str

type Trie struct {
	// 存储26个子节点，每个节点对应一个英文字母
	children [26]*Trie
	// 表示节点是否为字符串的结尾
	isEnd bool
}

func Constructor() Trie {
	return Trie{}
}

// 插入字符串
func (t *Trie) Insert(word string) {
	node := t
	for _, ch := range word {
		// ch为int32类型，'a'也为int32类型
		ch -= 'a'
		if node.children[ch] == nil {
			node.children[ch] = new(Trie)
		}
		// 向下伸展节点
		node = node.children[ch]
	}
	// 表示当前节点为字符串的结尾
	node.isEnd = true
}

// 搜索prefix字符串最后一个字符所在的节点
func (t *Trie) SearchPrefix(prefix string) *Trie {
	node := t
	for _, ch := range prefix {
		ch -= 'a'
		if node.children[ch] == nil {
			return nil
		}
		// 逐层往下查找
		node = node.children[ch]
	}
	return node
}

// 判断word字符串是否在前缀树中
func (t *Trie) Search(word string) bool {
	node := t.SearchPrefix(word)
	return node != nil && node.isEnd
}

// 判断prefix字符串是否在前缀树中
func (t *Trie) StartsWith(prefix string) bool {
	return t.SearchPrefix(prefix) != nil
}

/**
 * Your Trie object will be instantiated and called as such:
 * obj := Constructor();
 * obj.Insert(word);
 * param_2 := obj.Search(word);
 * param_3 := obj.StartsWith(prefix);
 */
